import argparse
import numpy as np
import handleDataAndDrawPicture
import networkx as nx
import os

def _findCommunity(edges):
    visit = [False for i in range(len(edges))]
    communities = []
    for i in range(len(visit)):
        if not visit[i]:
            community = [  ]
            temp = [i]
            newTemp = []
            while(len(temp)>0):
                for node in temp:
                    community.append(node)
                    visit[node] = True
                    for j in range(len(visit)):
                        if (edges[node,j] == 1) and (not visit[j]):
                            visit[j] = True
                            newTemp.append(j)
                temp = newTemp
                newTemp = []
            # community = [i + 1 for i in community]
            communities.append(community)
    return communities


def _havePath(edges, i, j):
    visit = [False for node in range(len(edges))]
    temp = [i]
    newTemp = []
    while(len(temp) > 0):
        for node in temp:
            visit[node] = True
            for nodej in range(len(visit)):
                if (edges[node,nodej] == 1) and (not visit[nodej]):
                    if nodej == j:
                        return True
                    visit[nodej] = True
                    newTemp.append(nodej)
            temp = newTemp
            newTemp = []
    return False

def _calQValue(edges):
    m = edges.sum()
    if m == 0:
        return -10
    Q = 0
    for nodei in range(len(edges)):
        for nodej in range(len(edges)):
            if _havePath(edges,nodej,nodei) :
                item = edges[nodei,nodej]-(edges[nodei].sum()*edges[nodej].sum()/(2*m))
                Q += item
    Q /= 2*m
    return Q

def _calBetweenness(edges):
    edgeBetweenness = np.zeros_like(edges, dtype=np.float64)
    for originNode in range(len(edges)): # 以originNode为原点出发计算edgeBetweeness
        # 做树
        visit = [False for i in range(len(edges))]
        path = []
        level = [originNode]
        newLevel = []
        visit[originNode] = True
        while len(level) > 0:
            for nodei in level :
                for nodej in range(len(edges)):
                    if (edges[nodei,nodej] == 1) and (not visit[nodej]):
                        newLevel.append(nodej)
                        visit[nodej] = True
            path.append(level)
            level = newLevel
            newLevel = []

        # 计算此种树状结构的w和d
        d = np.zeros(len(edges),dtype=np.int32)
        w = np.zeros(len(edges),dtype=np.int32)
        w[originNode] = 1
        for treeLevel in range(len(path) - 1):
            for nodeUp in path[treeLevel]:
                for nodeDown in path[treeLevel+1]:
                    if (edges[nodeUp,nodeDown] == 1 ) and (d[nodeDown] == 0) :
                        d[nodeDown] = d[nodeUp] + 1
                        w[nodeDown] = w[nodeUp]
                    elif(edges[nodeUp,nodeDown] == 1 ) and (d[nodeDown] != 0) and (d[nodeDown] == d[nodeUp] + 1):
                        w[nodeDown] += w[nodeUp]
        # 从下往上计算edge betweenness
        thisTreeEdgeBetweenness = np.zeros_like(edges, dtype=np.float64)
        nodeBetweenness = np.zeros(len(edges),dtype=np.float64)
        for treeLevel in range(len(path) - 1):
            for nodeDown in path[-1*(treeLevel+1)]:
                for nodeUp in path[-1*(treeLevel+2)]:
                    if edges[nodeUp,nodeDown] == 1:
                        thisTreeEdgeBetweenness[nodeUp,nodeDown] = (w[nodeUp] / w[nodeDown])*(1+nodeBetweenness[nodeDown])
                        thisTreeEdgeBetweenness[nodeDown,nodeUp] = thisTreeEdgeBetweenness[nodeUp,nodeDown]
                        nodeBetweenness[nodeUp] += thisTreeEdgeBetweenness[nodeUp,nodeDown]
        edgeBetweenness += thisTreeEdgeBetweenness
    return edgeBetweenness
        

def GNAlgorithm(originEdges):
    edges = originEdges.copy()
    while edges.sum() > 0:
        maxBetweenness = -5
        edgeBetweenness = _calBetweenness(edges)
        maxBetweennessNodei, maxBetweennessNodej = -1, -1
        for nodei in range(len(edges)):
            for nodej in range(nodei + 1, len(edges)):
                if edgeBetweenness[nodei,nodej] > maxBetweenness :
                    maxBetweenness = edgeBetweenness[nodei,nodej]
                    maxBetweennessNodei, maxBetweennessNodej = nodei, nodej
        edges[maxBetweennessNodei,maxBetweennessNodej] = 0
        edges[maxBetweennessNodej,maxBetweennessNodei] = 0
        if len(_findCommunity(edges)) == 2:
            return edges
        print(f'remove the edge between {maxBetweennessNodei} and {maxBetweennessNodej}')

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--data_path',default='./ucidata-zachary/out.ucidata-zachary',type=str,required=False,help='数据集路径')

    args = parser.parse_args()
    dataPath = args.data_path
    edges = handleDataAndDrawPicture.handleData(dataPath)
    bestEdges = GNAlgorithm(edges)
    communities = _findCommunity(bestEdges)
    print(f'\nthere are {len(communities)} communities')
    for i in range(len(communities)):
        print(f'No.{i} community is {communities[i]}')
    nodePos = handleDataAndDrawPicture.drawResult(communities,edges)
    handleDataAndDrawPicture.drawGraph(edges,nodePos)

if __name__  == '__main__':
    main()
